package com.wt.admin.service.ai.impl.mcp;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
import org.springframework.ai.mcp.client.autoconfigure.NamedClientMcpTransport;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpSseClientProperties;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpStdioClientProperties;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;

import java.net.http.HttpClient;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;


@Component
public class MCPTransport {

    ObjectMapper objectMapper = new ObjectMapper();


    /**
     * webFlux
     * @return
     */

    public List<NamedClientMcpTransport> webFluxClientTransports(Map<String, McpSseClientProperties.SseParameters> connections) {
        WebClient.Builder webClientBuilderTemplate = WebClient.builder();
        List<NamedClientMcpTransport> sseTransports = new ArrayList();
        connections.forEach((key, value) -> {
            WebClient.Builder webClientBuilder = webClientBuilderTemplate.clone().baseUrl(value.url());
            WebFluxSseClientTransport transport = new WebFluxSseClientTransport(webClientBuilder, objectMapper);
            sseTransports.add(new NamedClientMcpTransport(key, transport));
        });
        return sseTransports;
    }

    /**
     * sse 方式，重点传入 服务名称(name) 和 请求地址(url)
     * @return
     */
    public List<NamedClientMcpTransport> mcpHttpClientTransports(Map<String, McpSseClientProperties.SseParameters> connections) {
        List<NamedClientMcpTransport> sseTransports = new ArrayList();
        connections.forEach((key, value) -> {
            HttpClientSseClientTransport transport = HttpClientSseClientTransport.builder(value.url())
                    .sseEndpoint(value.sseEndpoint())
                    .clientBuilder(HttpClient.newBuilder())
                    .objectMapper(objectMapper)
                    .build();
            sseTransports.add(new NamedClientMcpTransport(key, transport));
        });
        return sseTransports;
    }

    public List<NamedClientMcpTransport> stdioTransports(String json) throws Exception {
        // json 转 Map<String, Map<String, Parameters>>
        ObjectMapper objectMapper = new ObjectMapper();
        Map<String, Map<String, McpStdioClientProperties.Parameters>> config = objectMapper.readValue(json,
                new TypeReference<Map<String, Map<String, McpStdioClientProperties.Parameters>>>() {});;

        List<NamedClientMcpTransport> stdoiTransports = new ArrayList<>();
        for (Map.Entry<String, McpStdioClientProperties.Parameters> entry : config.values().iterator().next().entrySet()) {
            String key = entry.getKey();
            McpStdioClientProperties.Parameters parameters = entry.getValue();
            StdioClientTransport transport = new StdioClientTransport(parameters.toServerParameters());
            stdoiTransports.add(new NamedClientMcpTransport(key, transport));
        }
        return stdoiTransports;
    }






}
